#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
@File        : mngs_mjn_select_pan.py
@Author      : Bing Liang
@Email       : believer19940901@gmail.com
@Date        : 2025/11/19 10:49
@Description :
    根据物种名，从泛基因组 FASTA 中提取该物种对应的序列，并按照参考基因组长度截取。
    输出到给定的 FASTA 文件。
"""

from argparse import ArgumentParser, Namespace
from pathlib import Path
from typing import Dict
from Bio import SeqIO


def _parse_pan_sp() -> Dict[str, str]:
    """
    获取物种与泛基因组文件名的映射关系。

    文件格式（pan_species.txt）示例：
        Bacillus cereus:>bacillus_cereus_pan.fasta

    Returns:
        Dict[str, str]: {species_name: pangenome_filename}
    """
    mapping = {}
    pan_sp_file = Path("/data/mNGS/pipeline/database/pangenome/pan_species.txt")

    with open(pan_sp_file, "r", encoding="utf-8") as f:
        for line in f:
            if line := line.strip():
                fasta_name, species_name = line.split(":>")
                mapping[species_name.strip()] = fasta_name.strip()

    return mapping


def _parse_len_sp() -> Dict[str, int]:
    """
    获取物种参考基因组长度，用于按长度截取泛基因组序列。

    文件格式（ref_length_4_coverage.xls）示例：
        Bacillus cereus\t5342910

    Returns:
        Dict[str, int]: {species_name: reference_length}
    """
    mapping = {}
    len_sp_file = Path("/data/mNGS/pipeline/script/table/ref_length_4_coverage.xls")

    with open(len_sp_file, "r", encoding="utf-8") as f:
        for line in f:
            if line := line.strip():
                species_name, length_str = line.split("\t")
                mapping[species_name.strip()] = int(length_str.strip())

    return mapping


def _parse_args() -> Namespace:
    """
    解析命令行参数。

    Returns:
        Namespace: 包含 species, out_file
    """
    parser = ArgumentParser(description="Select truncated pangenome sequence for a species.")
    parser.add_argument("--species", type=str, required=True, help="物种名称，与映射表一致。")
    parser.add_argument("--out_file", type=str, required=True, help="输出 FASTA 文件路径。")
    return parser.parse_args()


def main(args: Namespace) -> None:
    """
    主程序流程：
        1. 加载物种-泛基因组映射。
        2. 加载物种参考基因组长度。
        3. 从对应的泛基因组 FASTA 中找到同名序列。
        4. 按参考长度截断，写入输出文件。

    Args:
        args (Namespace): 命令行参数
    """
    species = args.species.strip()
    out_file = Path(args.out_file).absolute()
    out_file.parent.mkdir(parents=True, exist_ok=True)

    # 映射加载
    pan_sp_map = _parse_pan_sp()
    len_sp_map = _parse_len_sp()

    # 查找对应的泛基因组文件
    pangenome_file = Path("/data/mNGS/pipeline/database/pangenome") / pan_sp_map[species]

    ref_len = len_sp_map[species]

    # 查找 FASTA 并写出截断序列
    with open(pangenome_file, "r", encoding="utf-8") as fr, \
         open(out_file, "w", encoding="utf-8") as fw:

        for record in SeqIO.parse(fr, "fasta"):
            # 只导出与 species 名称匹配的序列
            if record.id == species:
                truncated_seq = record.seq[:ref_len]
                fw.write(f">{species}\n{truncated_seq}\n")
                break


if __name__ == "__main__":
    main(_parse_args())
